# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: chensili

import datetime

from recordclass import recordclass

from coin.base.book.types import TradeSide

BasicStats = recordclass('BasicStats',
                         [
                             'first_trade',
                             'last_trade',
                             'trade_count',
                             'avg_trade_size',
                             'each_hour_count',
                             'feed_checker_invalid_count'
                         ])

LongestNoTrade = recordclass('LongestNoTrade', ['begin', 'end', 'time_delta', 'last_timestamp'])

VolumeStats = recordclass('VolumeStats', ['total_volume', 'buy_volume', 'sell_volume'])

PriceStats = recordclass(
    'PriceStats',
    ['pct_return', 'price_open', 'price_high', 'price_low', 'price_close', 'vwap', 'twap'])

TradeStats = recordclass(
    'TradeStats', ['product', 'basic_stats', 'max_time_no_trade', 'volume_stats', 'price_stats'])


def gen_empty_TradeStats():
  return TradeStats(product=None,
                    basic_stats=BasicStats(None, None, 0, 0, [None] * 24, 0),
                    max_time_no_trade=LongestNoTrade(None, None, None, None),
                    volume_stats=VolumeStats(0, 0, 0),
                    price_stats=PriceStats(None, None, None, None, None, None, None))


class TradeStatsCalculator(object):
  def __init__(self, trade_stats):
    self._trade_stats = trade_stats
    self._prev_trade_timestamp = None
    self._total_time = 0

  def set_product(self, product):
    self.trade_stats.product = product

  @property
  def trade_stats(self):
    return self._trade_stats

  def update_by_trade(self, trade, feed_checker, product):
    self.update_basic_stats(trade)
    self.update_longest_no_trade(trade)
    self.update_volume_stats(trade)
    self.update_price_stats(trade)

    state = feed_checker.check_per_product_state(product=product, timestamp=trade.timestamp)
    if state.invalid:
      self.trade_stats.basic_stats.feed_checker_invalid_count += 1
      feed_checker.restart(timestamp=trade.timestamp)

  def update_basic_stats(self, trade):
    basic_stats = self.trade_stats.basic_stats
    if basic_stats.first_trade is None:
      basic_stats.first_trade = trade.timestamp
    if self._prev_trade_timestamp is None:
      self._prev_trade_timestamp = trade.timestamp
    else:
      self._prev_trade_timestamp = basic_stats.last_trade
    basic_stats.last_trade = trade.timestamp
    basic_stats.avg_trade_size = \
      (basic_stats.avg_trade_size * basic_stats.trade_count +
       trade.qty) / (basic_stats.trade_count + 1)
    basic_stats.trade_count += 1
    hour = datetime.datetime.fromtimestamp(trade.timestamp / (10.**9)).hour
    if basic_stats.each_hour_count[hour] is None:
      basic_stats.each_hour_count[hour] = 0
    basic_stats.each_hour_count[hour] += 1

  def update_longest_no_trade(self, trade):
    longest_no_trade = self.trade_stats.max_time_no_trade
    if longest_no_trade.begin is None:
      longest_no_trade.begin = trade.timestamp
      longest_no_trade.end = trade.timestamp
      longest_no_trade.time_delta = 0
      longest_no_trade.last_timestamp = trade.timestamp
    else:
      last_timestamp = longest_no_trade.last_timestamp
      curr_delta = trade.timestamp - last_timestamp
      if curr_delta > longest_no_trade.time_delta:
        longest_no_trade.begin = last_timestamp
        longest_no_trade.end = trade.timestamp
        longest_no_trade.time_delta = curr_delta
      longest_no_trade.last_timestamp = trade.timestamp

  def update_volume_stats(self, trade):
    volume_stats = self.trade_stats.volume_stats
    volume_stats.total_volume += trade.qty
    if trade.side == TradeSide.TRADE_BUY_SIDE:
      volume_stats.buy_volume += trade.qty
    elif trade.side == TradeSide.TRADE_SELL_SIDE:
      volume_stats.sell_volume += trade.qty
    else:
      assert trade.side == TradeSide.TRADE_NO_SIDE, \
        'Undefined order side: %s' % trade.side

  def update_price_stats(self, trade):
    price_stats = self.trade_stats.price_stats
    if price_stats.price_open is None:
      price_stats.price_open = trade.price
    price_stats.price_close = trade.price
    if price_stats.price_high is None:
      price_stats.price_high = trade.price
    else:
      if trade.price > price_stats.price_high:
        price_stats.price_high = trade.price
    if price_stats.price_low is None:
      price_stats.price_low = trade.price
    else:
      if trade.price < price_stats.price_low:
        price_stats.price_low = trade.price
    price_stats.pct_return = \
      price_stats.price_close / price_stats.price_open - 1

    curr_delta = (trade.timestamp - self._prev_trade_timestamp) / (10.**9)
    prev_total_time = self._total_time
    self._total_time += curr_delta
    if self._total_time > 0:
      if price_stats.twap is None:
        price_stats.twap = trade.price
      else:
        price_stats.twap = \
          (price_stats.twap * prev_total_time +
           trade.price * curr_delta) / self._total_time

    if price_stats.vwap is None:
      price_stats.vwap = trade.price
    else:
      total_volume = self.trade_stats.volume_stats.total_volume
      price_stats.vwap = \
        (price_stats.vwap * (total_volume - trade.qty) +
         trade.price * trade.qty) / total_volume
